import numpy as np
from PIL import Image


def pad_image(img, pad=10):
    """
    给图像四周添加白边（值为0，表示黑底图像中的白色字符）
    """
    return np.pad(img, ((pad, pad), (pad, pad)), constant_values=0)


def load_image(path):
    img = Image.open(path).convert("L")
    img = 255 - np.array(img)
    img = (img > 128).astype(int)
    return img

def split_rows(img, min_height=10):
    proj = np.sum(img, axis=1)
    in_row = False
    rows = []
    start = 0
    for i, val in enumerate(proj):
        if val > 0 and not in_row:
            in_row = True
            start = i
        elif val == 0 and in_row:
            in_row = False
            end = i
            if end - start > min_height:
                row_img = img[start:end, :]
                rows.append(row_img)
    return rows


def split_columns(img, min_width=10):
    proj = np.sum(img, axis=0)
    in_char = False
    chars = []
    start = 0
    for i, val in enumerate(proj):
        if val > 0 and not in_char:
            in_char = True
            start = i
        elif val == 0 and in_char:
            in_char = False
            end = i
            if end - start > min_width:
                char_img = img[:, start:end]
                chars.append(char_img)
    return chars


def normalize_char(img, size=(20, 20)):
    pil_img = Image.fromarray((img * 255).astype(np.uint8))
    pil_img = pil_img.resize(size, Image.NEAREST)
    return (np.array(pil_img) > 128).astype(int)


def build_templates_from_two_row_image(path):
    img = load_image(path)
    rows = split_rows(img)
    print(f"[DEBUG] 行数：{len(rows)}")  # 期望是 2

    templates = {}
    digits = list("12345") + list("67890")  # 对应顺序

    char_imgs = []
    for i, row in enumerate(rows):
        chars = split_columns(row)
        print(f"[DEBUG] 第{i + 1}行切分出字符数：{len(chars)}")  # 每行期望是 5
        char_imgs.extend(chars)

    print(f"[DEBUG] 总共字符数：{len(char_imgs)}")

    for i, digit in enumerate(digits):
        if i >= len(char_imgs):
            print(f"[ERROR] 第 {i} 个模板 '{digit}' 超出范围")
            break
        norm = normalize_char(char_imgs[i])
        v, h = np.sum(norm, axis=0), np.sum(norm, axis=1)
        templates[digit] = (v, h)
    return templates

# 识别函数
def recognize_char(char_img, templates):
    norm = normalize_char(char_img)
    v, h = np.sum(norm, axis=0), np.sum(norm, axis=1)
    return min(
        templates,
        key=lambda d: np.linalg.norm(v - templates[d][0])
        + np.linalg.norm(h - templates[d][1]),
    )


def recognize_digits_in_image(image_path, templates):
    img = load_image(image_path)
    img = pad_image(img, pad=10)

    rows = split_rows(img)
    all_chars = []

    for row in rows:
        chars = split_columns(row)
        all_chars.extend(chars)

    recognized = []
    for char_img in all_chars:
        digit = recognize_char(char_img, templates)
        recognized.append(digit)

    return recognized

if __name__ == "__main__":
    path = "./img.png"
    templates = build_templates_from_two_row_image(path)
    test_img = "./img_3.png"
    print("模板构建完成，包含数字：", list(templates.keys()))
    digits = recognize_digits_in_image(test_img, templates)
    print("识别结果：", digits)


